# 以下这段代码用于测试gsl最小二乘法计算各关节夹角
# 共涉及4个参考系: 机械臂基参考系（w），机械臂末端参考系（t），相机参考系（c），标记物参考系（a）
# 代码中用T_ij来表示由参考系（i）到参考系（j）的变换矩阵
# 代码中所有长度单位为 mm

import cv2
import numpy as np
import apriltag
import nllms
from functools import reduce
np.seterr(divide = "ignore", invalid = "ignore")

# 一、定义调用gsl最小二乘算法所需的三个函数

# 1.输入：各关节夹角joints
#         标记物3个角点在机械臂末端参考系下的齐次坐标f_t
#   输出：标记物3个角点在机械臂基参考系下的齐次坐标估计值f_est
#         雅克比矩阵J

def fk(joints, f_t):
    
    # 将各关节夹角用弧度值表示
    theta_1, theta_2, theta_3, theta_4, theta_5 = [joints[i]/180*np.pi for i in range(5)]
    
    # DH参数   
    DH = [[theta_1, 0, 8, np.pi/2],
          [theta_2, 232, 34, np.pi/2],
          [theta_3, 0, 385, 0],
          [theta_4, 0, 82, -np.pi/2],
          [theta_5, 367, 0, 0]]
    
    T_1 = []
    T_2 = []
    
    for i in range(5):
        theta, d, a, alpha = DH[i]
        
        # 简化三角函数计算
        cth, sth = np.cos(theta), np.sin(theta)
        cal, sal = np.cos(alpha), np.sin(alpha)
        
        # 相邻关节参考系间的变换矩阵
        t_1 = np.array([[cth, sth, 0, a],
                        [sth*cal, cth*cal, -sal, -sal*d],
                        [sth*sal, cth*sal, cal, cal*d],
                        [0, 0, 0, 1]])
        T_1.append(t_1)
        
        # 关节i与关节i+1之间的变换矩阵对关节i+1的夹角求导
        t_2 = np.array([[-sth, cth, 0, 0],
                        [cth*cal, -sth*cal, 0, 0],
                        [cth*sal, -sth*cal, 0, 0],
                        [0, 0, 0, 1]])
        T_2.append(t_2)
    # 根据正向运动学，所有相邻关节参考系间的变换矩阵按顺序依次相乘，得到首尾两个关节参考系间的变换矩阵
    T_wt = reduce(np.dot, T_1)
    f_est = T_wt.dot(f_t)
    
    # f_est分别对theta_1、theta_2、theta_3、theta_4、theta_5求偏导
    J_1 = T_2[0].dot(T_1[1]).dot(T_1[2]).dot(T_1[3]).dot(T_1[4]).dot(f_t)
    J_2 = T_1[0].dot(T_2[1]).dot(T_1[2]).dot(T_1[3]).dot(T_1[4]).dot(f_t)
    J_3 = T_1[0].dot(T_1[1]).dot(T_2[2]).dot(T_1[3]).dot(T_1[4]).dot(f_t)
    J_4 = T_1[0].dot(T_1[1]).dot(T_1[2]).dot(T_2[3]).dot(T_1[4]).dot(f_t)
    J_5 = T_1[0].dot(T_1[1]).dot(T_1[2]).dot(T_1[3]).dot(T_2[4]).dot(f_t)
    
    # 按照规定的排列方式组合雅克比矩阵
    J_1_ = [i for i in J_1[0:3].flat]
    J_2_ = [i for i in J_2[0:3].flat]
    J_3_ = [i for i in J_3[0:3].flat]
    J_4_ = [i for i in J_4[0:3].flat]
    J_5_ = [i for i in J_5[0:3].flat]
    J = np.vstack((J_1_, J_2_, J_3_, J_4_, J_5_))
    
    # 组合的雅克比矩阵为（5，9）矩阵，还需进行一次转置
    J = np.transpose(J)
    
    return (f_est, J)
    
# 2.输入：环境参数env
#         自变量xs(各关节夹角)
#   输出：目标函数Fs = 估计值f_est - 实测值f_w

def fCb_(env, xs):
    (f_w, f_t) = env
    f_est, jac_est = fk(xs, f_t)
    # f_est和f_tar为(4, 3)矩阵，所需的目标函数为(1, 9)向量
    Fs = [v1-v2 for v1, v2 in zip(f_est[:3].flat, f_w[:3].flat)]
    return Fs

# 3.输入：环境参数env
#         自变量xs(各关节夹角)
#   输出：雅可比矩阵jac_est

def jacCb_(env, xs):
    (f_w, f_t) = env
    f_est, jac_est = fk(xs, f_t)
# 得到的jac_est为（9，5）矩阵
    return jac_est


# 二、人工测量参数
# 相机内参
camera_matrix = np.array(([2054, 0, 631],
                         [0, 2051, 422],
                         [0, 0, 1]), dtype = np.double)

# 相机参考系(c)到机械臂末端参考系(t)的变换矩阵
T_ct = np.array([[1, 0, 0, 0],
                 [0, 1, 0, -87],
                 [0, 0, 1, 243],
                 [0, 0, 0, 1]])

# 机械臂基参考系(w)到标记物参考系(a)的变换矩阵
T_wa = np.array([[1, 0, 0, -61],
                 [0, 1, 0, 105],
                 [0, 0, 1, 1885],
                 [0, 0, 0, 1]])

# 标记物3个角点在标记物参考系下的齐次坐标
f_a = [[-20, 20, -20],
       [20, 20, -20],
       [0, 0, 0],
       [1, 1, 1]]

# 在机械臂基参考系下,标记物三个角点齐次坐标的实测值
f_w = T_wa.dot(f_a)


# 三、视觉识别标记物与相机的相对姿态
# 读取相机拍摄的图像，对图像进行预处理
img = cv2.imread("1.jpg")  #文档中出现的两张测试图像已上传，读取图像时麻烦改一下名字
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

# 识别apriltag码四个角点的像素坐标：
tag_detector = apriltag.Detector()
tags = tag_detector.detect(gray)
# 共识别出大小两个apriltag码，记录较小的apriltag码的四个角点像素坐标
points_2d = np.array((tags[1].corners[0],                           # 左下
                      tags[1].corners[1],                           # 左上
                      tags[1].corners[2],                           # 右上
                      tags[1].corners[3]), dtype = np.double)       # 右下

# apriltag码边长为4cm，标记物参考系以apriltag码中心作为坐标原点，且apriltag四个角点位于z=0平面上
# apriltag码四个角点在标记物参考系下的坐标：
points_3d = np.array(([-20, 20, 0], [20, 20, 0], [20, -20, 0], [-20, -20, 0]), dtype = np.double)

# 运用PnP算法得到标记物参考系（a）到相机参考系（c）的变换矩阵
# 假定相机无畸变
dist_coefs = np.array([0, 0, 0, 0, 0], dtype = np.double)
retval, rvec, tvec = cv2.solvePnP(points_3d, points_2d, camera_matrix, dist_coefs)

# 将求解出来的旋转矢量以旋转矩阵的形式表示：
rotM = cv2.Rodrigues(rvec)[0]

# 将位移矢量和旋转矩阵组合成变换矩阵
I = np.array([[0, 0, 0, 1]])
Rt = np.hstack((rotM, tvec))
T_ac = np.vstack((Rt, I))

# 标记物三个角点在机械臂末端参考系下的坐标
T_ca = np.linalg.inv(T_ac)
f_t = T_tc.dot(T_ca).dot(f_a)


# 四、最小二乘计算
# 用python建立自变量（关节夹角）对象，并设置迭代初值
xs = [90, 0, 90, 80, 10]

# 用C语言建立自变量对象，并赋值
x = nllms.ffi.new("double []", 5)
for i, j in eumerate(xs):
    x[i] = j

# 将函数参数与前面定义的函数fCb_、jacCb_作为一个元组打包在一起作为一个整体“环境参数”传入
# 其中“函数参数”指调用fCb_、jacCb_时所需传入的参数
paramsCb = nllms.ffi.new_handle(((f_w, f_t), fCb_, jacCb_))

# 调用gsl最小二乘算法进行迭代计算。ret_为返回值，非零表示出错
ret_ = nllms.lib.search_by_nonlinear_lms(9, 5, paramsCb, 1.e-4, x)
print("search return {ret_}".format(ret_ = ret_,))

# 输出最终计算结果
joints_cur = [x[i] for i in range(5)]
print(joints_cur)

